{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "AlphaFold2_advanced.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/AlphaFold2_advanced.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pc5-mbsX9PZC"
      },
      "source": [
        "# AlphaFold2_advanced\n",
        "\n",
        "-----------------\n",
        "- <b><font color='green'>21Aug2021: MMseqs2 API has finished upgrade, all should be ready to go! Report any errors.</font></b>\n",
        "\n",
        "-----------------\n",
        "\n",
        "This notebook modifies deepmind's [original notebook](https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb) to add experimental support for modeling complexes (both homo and hetero-oligomers), option to run MMseqs2 instead of Jackhmmer for MSA generation and advanced functionality.\n",
        "\n",
        "See [ColabFold](https://github.com/sokrypton/ColabFold/) for other related notebooks\n",
        "\n",
        "**Limitations**\n",
        "- This notebook does NOT use Templates.\n",
        "- For a typical Google-Colab session, with a `16G-GPU`, the max total length is **1400 residues**. Sometimes a `12G-GPU` is assigned, in which the max length is ~1000 residues.\n",
        "- Can I use the models for **Molecular Replacement**? Yes, but be CAREFUL, the bfactor column is populated with pLDDT confidence values (higher = better). Phenix.phaser expects a \"real\" bfactor, where (lower = better). See [post](https://twitter.com/cheshireminima/status/1423929241675120643) from Claudia Millán on how to process models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "woIxeCPygt7K",
        "cellView": "form"
      },
      "source": [
        "#@title Install software\n",
        "#@markdown Please execute this cell by pressing the _Play_ button \n",
        "#@markdown on the left.\n",
        "\n",
        "import os\n",
        "import tensorflow as tf\n",
        "tf.config.set_visible_devices([], 'GPU')\n",
        "\n",
        "import jax\n",
        "if jax.local_devices()[0].platform == 'tpu':\n",
        "  raise RuntimeError('Colab TPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "elif jax.local_devices()[0].platform == 'cpu':\n",
        "  raise RuntimeError('Colab CPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "\n",
        "from IPython.utils import io\n",
        "import subprocess\n",
        "import tqdm.notebook\n",
        "\n",
        "from sys import version_info \n",
        "python_version = f\"{version_info.major}.{version_info.minor}\"\n",
        "\n",
        "GIT_REPO = 'https://github.com/deepmind/alphafold'\n",
        "SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar'\n",
        "PARAMS_DIR = './alphafold/data/params'\n",
        "PARAMS_PATH = os.path.join(PARAMS_DIR, os.path.basename(SOURCE_URL))\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "\n",
        "# if not already installed\n",
        "try:\n",
        "  total = 55\n",
        "  with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "    with io.capture_output() as captured:\n",
        "      if not os.path.isdir(\"alphafold\"):\n",
        "        %shell rm -rf alphafold\n",
        "        %shell git clone {GIT_REPO} alphafold\n",
        "        %shell (cd alphafold; git checkout 1d43aaff941c84dc56311076b58795797e49107b --quiet)\n",
        "\n",
        "        # colabfold patches\n",
        "        %shell mkdir --parents tmp\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/colabfold.py\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/pairmsa.py\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/protein.patch -P tmp/\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/config.patch -P tmp/\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/model.patch -P tmp/\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/modules.patch -P tmp/\n",
        "\n",
        "        # install hhsuite\n",
        "        %shell curl -fsSL https://github.com/soedinglab/hh-suite/releases/download/v3.3.0/hhsuite-3.3.0-SSE2-Linux.tar.gz | tar xz -C tmp/\n",
        "\n",
        "        # Apply multi-chain patch from Lim Heo @huhlim\n",
        "        %shell patch -u alphafold/alphafold/common/protein.py -i tmp/protein.patch\n",
        "        \n",
        "        # Apply patch to dynamically control number of recycles (idea from Ryan Kibler)\n",
        "        %shell patch -u alphafold/alphafold/model/model.py -i tmp/model.patch\n",
        "        %shell patch -u alphafold/alphafold/model/modules.py -i tmp/modules.patch\n",
        "        %shell patch -u alphafold/alphafold/model/config.py -i tmp/config.patch\n",
        "        pbar.update(4)\n",
        "\n",
        "        %shell pip3 install biopython dm-haiku ml-collections py3Dmol\n",
        "        %shell pip3 install --no-dependencies ./alphafold\n",
        "        pbar.update(6)\n",
        "\n",
        "        # speedup from kaczmarj\n",
        "        %shell mkdir --parents \"{PARAMS_DIR}\"\n",
        "        %shell curl -fsSL \"{SOURCE_URL}\" | tar x -C \"{PARAMS_DIR}\"\n",
        "        pbar.update(14+27)\n",
        "\n",
        "        #######################################################################\n",
        "        %shell sudo apt install --quiet --yes hmmer\n",
        "        pbar.update(3)\n",
        "\n",
        "        # Create a ramdisk to store a database chunk to make Jackhmmer run fast.\n",
        "        %shell sudo mkdir -m 777 --parents /tmp/ramdisk\n",
        "        %shell sudo mount -t tmpfs -o size=9G ramdisk /tmp/ramdisk\n",
        "        pbar.update(1)\n",
        "      else:\n",
        "        pbar.update(55)\n",
        "\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise\n",
        "\n",
        "########################################################################################\n",
        "# --- Python imports ---\n",
        "import colabfold as cf\n",
        "import pairmsa\n",
        "import sys\n",
        "import pickle\n",
        "\n",
        "if \"/content/tmp/bin\" not in os.environ['PATH']:\n",
        "  os.environ['PATH'] += \":/content/tmp/bin:/content/tmp/scripts\"\n",
        "\n",
        "from urllib import request\n",
        "from concurrent import futures\n",
        "from google.colab import files\n",
        "import json\n",
        "from matplotlib import gridspec\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import py3Dmol\n",
        "\n",
        "from alphafold.model import model\n",
        "from alphafold.model import config\n",
        "from alphafold.model import data\n",
        "\n",
        "from alphafold.data import parsers\n",
        "from alphafold.data import pipeline\n",
        "from alphafold.data.tools import jackhmmer\n",
        "\n",
        "from alphafold.common import protein\n",
        "\n",
        "def run_jackhmmer(sequence, prefix):\n",
        "\n",
        "  fasta_path = f\"{prefix}.fasta\"\n",
        "  with open(fasta_path, 'wt') as f:\n",
        "    f.write(f'>query\\n{sequence}')\n",
        "\n",
        "  pickled_msa_path = f\"{prefix}.jackhmmer.pickle\"\n",
        "  if os.path.isfile(pickled_msa_path):\n",
        "    msas_dict = pickle.load(open(pickled_msa_path,\"rb\"))\n",
        "    msas, deletion_matrices, names = (msas_dict[k] for k in ['msas', 'deletion_matrices', 'names'])\n",
        "    full_msa = []\n",
        "    for msa in msas:\n",
        "      full_msa += msa\n",
        "  else:\n",
        "    # --- Find the closest source ---\n",
        "    test_url_pattern = 'https://storage.googleapis.com/alphafold-colab{:s}/latest/uniref90_2021_03.fasta.1'\n",
        "    ex = futures.ThreadPoolExecutor(3)\n",
        "    def fetch(source):\n",
        "      request.urlretrieve(test_url_pattern.format(source))\n",
        "      return source\n",
        "    fs = [ex.submit(fetch, source) for source in ['', '-europe', '-asia']]\n",
        "    source = None\n",
        "    for f in futures.as_completed(fs):\n",
        "      source = f.result()\n",
        "      ex.shutdown()\n",
        "      break\n",
        "\n",
        "    jackhmmer_binary_path = '/usr/bin/jackhmmer'\n",
        "    dbs = []\n",
        "\n",
        "    num_jackhmmer_chunks = {'uniref90': 59, 'smallbfd': 17, 'mgnify': 71}\n",
        "    total_jackhmmer_chunks = sum(num_jackhmmer_chunks.values())\n",
        "    with tqdm.notebook.tqdm(total=total_jackhmmer_chunks, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      def jackhmmer_chunk_callback(i):\n",
        "        pbar.update(n=1)\n",
        "\n",
        "      pbar.set_description('Searching uniref90')\n",
        "      jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(\n",
        "          binary_path=jackhmmer_binary_path,\n",
        "          database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/uniref90_2021_03.fasta',\n",
        "          get_tblout=True,\n",
        "          num_streamed_chunks=num_jackhmmer_chunks['uniref90'],\n",
        "          streaming_callback=jackhmmer_chunk_callback,\n",
        "          z_value=135301051)\n",
        "      dbs.append(('uniref90', jackhmmer_uniref90_runner.query(fasta_path)))\n",
        "\n",
        "      pbar.set_description('Searching smallbfd')\n",
        "      jackhmmer_smallbfd_runner = jackhmmer.Jackhmmer(\n",
        "          binary_path=jackhmmer_binary_path,\n",
        "          database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/bfd-first_non_consensus_sequences.fasta',\n",
        "          get_tblout=True,\n",
        "          num_streamed_chunks=num_jackhmmer_chunks['smallbfd'],\n",
        "          streaming_callback=jackhmmer_chunk_callback,\n",
        "          z_value=65984053)\n",
        "      dbs.append(('smallbfd', jackhmmer_smallbfd_runner.query(fasta_path)))\n",
        "\n",
        "      pbar.set_description('Searching mgnify')\n",
        "      jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(\n",
        "          binary_path=jackhmmer_binary_path,\n",
        "          database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/mgy_clusters_2019_05.fasta',\n",
        "          get_tblout=True,\n",
        "          num_streamed_chunks=num_jackhmmer_chunks['mgnify'],\n",
        "          streaming_callback=jackhmmer_chunk_callback,\n",
        "          z_value=304820129)\n",
        "      dbs.append(('mgnify', jackhmmer_mgnify_runner.query(fasta_path)))\n",
        "\n",
        "    # --- Extract the MSAs and visualize ---\n",
        "    # Extract the MSAs from the Stockholm files.\n",
        "    # NB: deduplication happens later in pipeline.make_msa_features.\n",
        "\n",
        "    mgnify_max_hits = 501\n",
        "    msas = []\n",
        "    deletion_matrices = []\n",
        "    names = []\n",
        "    for db_name, db_results in dbs:\n",
        "      unsorted_results = []\n",
        "      for i, result in enumerate(db_results):\n",
        "        msa, deletion_matrix, target_names = parsers.parse_stockholm(result['sto'])\n",
        "        e_values_dict = parsers.parse_e_values_from_tblout(result['tbl'])\n",
        "        e_values = [e_values_dict[t.split('/')[0]] for t in target_names]\n",
        "        zipped_results = zip(msa, deletion_matrix, target_names, e_values)\n",
        "        if i != 0:\n",
        "          # Only take query from the first chunk\n",
        "          zipped_results = [x for x in zipped_results if x[2] != 'query']\n",
        "        unsorted_results.extend(zipped_results)\n",
        "      sorted_by_evalue = sorted(unsorted_results, key=lambda x: x[3])\n",
        "      db_msas, db_deletion_matrices, db_names, _ = zip(*sorted_by_evalue)\n",
        "      if db_msas:\n",
        "        if db_name == 'mgnify':\n",
        "          db_msas = db_msas[:mgnify_max_hits]\n",
        "          db_deletion_matrices = db_deletion_matrices[:mgnify_max_hits]\n",
        "          db_names = db_names[:mgnify_max_hits]\n",
        "        msas.append(db_msas)\n",
        "        deletion_matrices.append(db_deletion_matrices)\n",
        "        names.append(db_names)\n",
        "        msa_size = len(set(db_msas))\n",
        "        print(f'{msa_size} Sequences Found in {db_name}')\n",
        "\n",
        "      pickle.dump({\"msas\":msas,\n",
        "                   \"deletion_matrices\":deletion_matrices,\n",
        "                   \"names\":names}, open(pickled_msa_path,\"wb\"))\n",
        "  return msas, deletion_matrices, names"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rowN0bVYLe9n",
        "cellView": "form"
      },
      "source": [
        "#@title Enter the amino acid sequence to fold ⬇️\n",
        "import re\n",
        "\n",
        "# define sequence\n",
        "sequence = 'PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK' #@param {type:\"string\"}\n",
        "sequence = re.sub(\"[^A-Z:/]\", \"\", sequence.upper())\n",
        "sequence = re.sub(\":+\",\":\",sequence)\n",
        "sequence = re.sub(\"/+\",\"/\",sequence)\n",
        "sequence = re.sub(\"^[:/]+\",\"\",sequence)\n",
        "sequence = re.sub(\"[:/]+$\",\"\",sequence)\n",
        "\n",
        "jobname = \"test\" #@param {type:\"string\"}\n",
        "jobname = re.sub(r'\\W+', '', jobname)\n",
        "\n",
        "# define number of copies\n",
        "homooligomer =  \"1\" #@param {type:\"string\"}\n",
        "homooligomer = re.sub(\"[:/]+\",\":\",homooligomer)\n",
        "homooligomer = re.sub(\"^[:/]+\",\"\",homooligomer)\n",
        "homooligomer = re.sub(\"[:/]+$\",\"\",homooligomer)\n",
        "\n",
        "if len(homooligomer) == 0: homooligomer = \"1\"\n",
        "homooligomer = re.sub(\"[^0-9:]\", \"\", homooligomer)\n",
        "homooligomers = [int(h) for h in homooligomer.split(\":\")]\n",
        "\n",
        "#@markdown - `sequence` Specify protein sequence to be modelled.\n",
        "#@markdown  - Use `/` to specify intra-protein chainbreaks (for trimming regions within protein).\n",
        "#@markdown  - Use `:` to specify inter-protein chainbreaks (for modeling protein-protein hetero-complexes).\n",
        "#@markdown  - For example, sequence `AC/DE:FGH` will be modelled as polypeptides: `AC`, `DE` and `FGH`. A seperate MSA will be generates for `ACDE` and `FGH`.\n",
        "#@markdown    If `pair_msa` is enabled, `ACDE`'s MSA will be paired with `FGH`'s MSA.\n",
        "#@markdown - `homooligomer` Define number of copies in a homo-oligomeric assembly.\n",
        "#@markdown  - Use `:` to specify different homooligomeric state (copy numer) for each component of the complex. \n",
        "#@markdown  - For example, **sequence:**`ABC:DEF`, **homooligomer:** `2:1`, the first protein `ABC` will be modeled as a homodimer (2 copies) and second `DEF` a monomer (1 copy).\n",
        "\n",
        "ori_sequence = sequence\n",
        "sequence = sequence.replace(\"/\",\"\").replace(\":\",\"\")\n",
        "seqs = ori_sequence.replace(\"/\",\"\").split(\":\")\n",
        "\n",
        "if len(seqs) != len(homooligomers):\n",
        "  if len(homooligomers) == 1:\n",
        "    homooligomers = [homooligomers[0]] * len(seqs)\n",
        "    homooligomer = \":\".join([str(h) for h in homooligomers])\n",
        "  else:\n",
        "    while len(seqs) > len(homooligomers):\n",
        "      homooligomers.append(1)\n",
        "    homooligomers = homooligomers[:len(seqs)]\n",
        "    homooligomer = \":\".join([str(h) for h in homooligomers])\n",
        "    print(\"WARNING: Mismatch between number of breaks ':' in 'sequence' and 'homooligomer' definition\")\n",
        "\n",
        "full_sequence = \"\".join([s*h for s,h in zip(seqs,homooligomers)])\n",
        "\n",
        "# prediction directory\n",
        "output_dir = 'prediction_' + jobname + '_' + cf.get_hash(full_sequence)[:5]\n",
        "os.makedirs(output_dir, exist_ok=True)\n",
        "# delete existing files in working directory\n",
        "for f in os.listdir(output_dir):\n",
        "  os.remove(os.path.join(output_dir, f))\n",
        "\n",
        "MIN_SEQUENCE_LENGTH = 16\n",
        "MAX_SEQUENCE_LENGTH = 2500\n",
        "\n",
        "aatypes = set('ACDEFGHIKLMNPQRSTVWY')  # 20 standard aatypes\n",
        "if not set(full_sequence).issubset(aatypes):\n",
        "  raise Exception(f'Input sequence contains non-amino acid letters: {set(sequence) - aatypes}. AlphaFold only supports 20 standard amino acids as inputs.')\n",
        "if len(full_sequence) < MIN_SEQUENCE_LENGTH:\n",
        "  raise Exception(f'Input sequence is too short: {len(full_sequence)} amino acids, while the minimum is {MIN_SEQUENCE_LENGTH}')\n",
        "if len(full_sequence) > MAX_SEQUENCE_LENGTH:\n",
        "  raise Exception(f'Input sequence is too long: {len(full_sequence)} amino acids, while the maximum is {MAX_SEQUENCE_LENGTH}. Please use the full AlphaFold system for long sequences.')\n",
        "\n",
        "if len(full_sequence) > 1400:\n",
        "  print(f\"WARNING: For a typical Google-Colab-GPU (16G) session, the max total length is ~1400 residues. You are at {len(full_sequence)}! Run Alphafold may crash.\")\n",
        "\n",
        "print(f\"homooligomer: '{homooligomer}'\")\n",
        "print(f\"total_length: '{len(full_sequence)}'\")\n",
        "print(f\"working_directory: '{output_dir}'\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ITcPnLkLuDDE",
        "cellView": "form"
      },
      "source": [
        "#@title Search against genetic databases\n",
        "#@markdown Once this cell has been executed, you will see\n",
        "#@markdown statistics about the multiple sequence alignment \n",
        "#@markdown (MSA) that will be used by AlphaFold. In particular, \n",
        "#@markdown you’ll see how well each residue is covered by similar \n",
        "#@markdown sequences in the MSA.\n",
        "#@markdown (Note that the search against databases and the actual prediction can take some time, from minutes to hours, depending on the length of the protein and what type of GPU you are allocated by Colab.)\n",
        "\n",
        "#@markdown ---\n",
        "msa_method = \"mmseqs2\" #@param [\"mmseqs2\",\"jackhmmer\",\"single_sequence\",\"precomputed\"]\n",
        "#@markdown - `mmseqs2` - FAST method from [ColabFold](https://github.com/sokrypton/ColabFold)\n",
        "#@markdown - `jackhmmer` - default method from Deepmind (SLOW, but may find more/less sequences).\n",
        "#@markdown - `single_sequence` - use single sequence input\n",
        "#@markdown - `precomputed` If you have previously run this notebook and saved the results,\n",
        "#@markdown you can skip this step by uploading \n",
        "#@markdown the previously generated  `prediction_?????/msa.pickle`\n",
        "\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown **custom msa options**\n",
        "add_custom_msa = False #@param {type:\"boolean\"}\n",
        "msa_format = \"fas\" #@param [\"fas\",\"a2m\",\"a3m\",\"sto\",\"psi\",\"clu\"]\n",
        "#@markdown - `add_custom_msa` - If enabled, you'll get an option to upload your custom MSA in the specified `msa_format`. Note: Your MSA will be supplemented with those from 'mmseqs2' or 'jackhmmer', unless `msa_method` is set to 'single_sequence'.\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown **pair msa options**\n",
        "\n",
        "#@markdown Experimental option for protein complexes. Pairing currently only supported for proteins in same operon (prokaryotic genomes).\n",
        "pair_mode = \"unpaired\" #@param [\"unpaired\",\"unpaired+paired\",\"paired\"] {type:\"string\"}\n",
        "#@markdown - `unpaired` - generate seperate MSA for each protein.\n",
        "#@markdown - `unpaired+paired` - attempt to pair sequences from the same operon within the genome. \n",
        "#@markdown - `paired` - only use sequences that were sucessfully paired.\n",
        "\n",
        "#@markdown Options to prefilter each MSA before pairing. (It might help if there are any paralogs in the complex.)\n",
        "pair_cov = 50 #@param [0,25,50,75,90] {type:\"raw\"}\n",
        "pair_qid = 20 #@param [0,15,20,30,40,50] {type:\"raw\"}\n",
        "#@markdown - `pair_cov` prefilter each MSA to minimum coverage with query (%) before pairing.\n",
        "#@markdown - `pair_qid` prefilter each MSA to minimum sequence identity with query (%) before pairing.\n",
        "\n",
        "# --- Search against genetic databases ---\n",
        "os.makedirs('tmp', exist_ok=True)\n",
        "msas, deletion_matrices = [],[]\n",
        "\n",
        "if add_custom_msa:\n",
        "  print(f\"upload custom msa in '{msa_format}' format\")\n",
        "  msa_dict = files.upload()\n",
        "  lines = msa_dict[list(msa_dict.keys())[0]].decode()\n",
        "  \n",
        "  # convert to a3m\n",
        "  with open(f\"tmp/upload.{msa_format}\",\"w\") as tmp_upload:\n",
        "    tmp_upload.write(lines)\n",
        "  os.system(f\"reformat.pl {msa_format} a3m tmp/upload.{msa_format} tmp/upload.a3m\")\n",
        "  a3m_lines = open(\"tmp/upload.a3m\",\"r\").read()\n",
        "\n",
        "  # parse\n",
        "  msa, mtx = parsers.parse_a3m(a3m_lines)\n",
        "  msas.append(msa)\n",
        "  deletion_matrices.append(mtx)\n",
        "\n",
        "  if len(msas[0][0]) != len(sequence):\n",
        "    raise ValueError(\"ERROR: the length of msa does not match input sequence\")\n",
        "\n",
        "if msa_method == \"precomputed\":\n",
        "  print(\"upload precomputed pickled msa from previous run\")\n",
        "  pickled_msa_dict = files.upload()\n",
        "  msas_dict = pickle.loads(pickled_msa_dict[list(pickled_msa_dict.keys())[0]])\n",
        "  msas, deletion_matrices = (msas_dict[k] for k in ['msas', 'deletion_matrices'])\n",
        "\n",
        "elif msa_method == \"single_sequence\":\n",
        "  if len(msas) == 0:\n",
        "    msas.append([sequence])\n",
        "    deletion_matrices.append([[0]*len(sequence)])\n",
        "\n",
        "else:\n",
        "  seqs = ori_sequence.replace('/','').split(':')\n",
        "  _blank_seq = [\"-\" * len(seq) for seq in seqs]\n",
        "  _blank_mtx = [[0] * len(seq) for seq in seqs]\n",
        "  def _pad(ns,vals,mode):\n",
        "    if mode == \"seq\": _blank = _blank_seq.copy()\n",
        "    if mode == \"mtx\": _blank = _blank_mtx.copy()\n",
        "    if isinstance(ns, list):\n",
        "      for n,val in zip(ns,vals): _blank[n] = val\n",
        "    else: _blank[ns] = vals\n",
        "    if mode == \"seq\": return \"\".join(_blank)\n",
        "    if mode == \"mtx\": return sum(_blank,[])\n",
        "\n",
        "  if len(seqs) == 1 or \"unpaired\" in pair_mode:\n",
        "    # gather msas\n",
        "    if msa_method == \"mmseqs2\":\n",
        "      prefix = cf.get_hash(\"\".join(seqs))\n",
        "      prefix = os.path.join('tmp',prefix)\n",
        "      print(f\"running mmseqs2\")\n",
        "      A3M_LINES = cf.run_mmseqs2(seqs, prefix, filter=True)\n",
        "\n",
        "    for n, seq in enumerate(seqs):\n",
        "      # tmp directory\n",
        "      prefix = cf.get_hash(seq)\n",
        "      prefix = os.path.join('tmp',prefix)\n",
        "\n",
        "      if msa_method == \"mmseqs2\":\n",
        "        # run mmseqs2\n",
        "        a3m_lines = A3M_LINES[n]\n",
        "        msa, mtx = parsers.parse_a3m(a3m_lines)\n",
        "        msas_, mtxs_ = [msa],[mtx]\n",
        "\n",
        "      elif msa_method == \"jackhmmer\":\n",
        "        print(f\"running jackhmmer on seq_{n}\")\n",
        "        # run jackhmmer\n",
        "        msas_, mtxs_, names_ = ([sum(x,())] for x in run_jackhmmer(seq, prefix))\n",
        "      \n",
        "      # pad sequences\n",
        "      for msa_,mtx_ in zip(msas_,mtxs_):\n",
        "        msa,mtx = [sequence],[[0]*len(sequence)]      \n",
        "        for s,m in zip(msa_,mtx_):\n",
        "          msa.append(_pad(n,s,\"seq\"))\n",
        "          mtx.append(_pad(n,m,\"mtx\"))\n",
        "\n",
        "        msas.append(msa)\n",
        "        deletion_matrices.append(mtx)\n",
        "\n",
        "  ####################################################################################\n",
        "  # PAIR_MSA\n",
        "  ####################################################################################\n",
        "  if len(seqs) > 1 and (pair_mode == \"paired\" or pair_mode == \"unpaired+paired\"):\n",
        "    print(\"attempting to pair some sequences...\")\n",
        "\n",
        "    if msa_method == \"mmseqs2\":\n",
        "      prefix = cf.get_hash(\"\".join(seqs))\n",
        "      prefix = os.path.join('tmp',prefix)\n",
        "      print(f\"running mmseqs2_noenv_nofilter on all seqs\")\n",
        "      A3M_LINES = cf.run_mmseqs2(seqs, prefix, use_env=False, use_filter=False)\n",
        "\n",
        "    _data = []\n",
        "    for a in range(len(seqs)):\n",
        "      print(f\"prepping seq_{a}\")\n",
        "      _seq = seqs[a]\n",
        "      _prefix = os.path.join('tmp',cf.get_hash(_seq))\n",
        "\n",
        "      if msa_method == \"mmseqs2\":\n",
        "        a3m_lines = A3M_LINES[a]\n",
        "        _msa, _mtx, _lab = pairmsa.parse_a3m(a3m_lines,\n",
        "                                             filter_qid=pair_qid/100,\n",
        "                                             filter_cov=pair_cov/100)\n",
        "\n",
        "      elif msa_method == \"jackhmmer\":\n",
        "        _msas, _mtxs, _names = run_jackhmmer(_seq, _prefix)\n",
        "        _msa, _mtx, _lab = pairmsa.get_uni_jackhmmer(_msas[0], _mtxs[0], _names[0],\n",
        "                                                     filter_qid=pair_qid/100,\n",
        "                                                     filter_cov=pair_cov/100)\n",
        "\n",
        "      if len(_msa) > 1:\n",
        "        _data.append(pairmsa.hash_it(_msa, _lab, _mtx, call_uniprot=False))\n",
        "      else:\n",
        "        _data.append(None)\n",
        "    \n",
        "    Ln = len(seqs)\n",
        "    O = [[None for _ in seqs] for _ in seqs]\n",
        "    for a in range(Ln):\n",
        "      if _data[a] is not None:\n",
        "        for b in range(a+1,Ln):\n",
        "          if _data[b] is not None:\n",
        "            print(f\"attempting pairwise stitch for {a} {b}\")            \n",
        "            O[a][b] = pairmsa._stitch(_data[a],_data[b])\n",
        "            _seq_a, _seq_b, _mtx_a, _mtx_b = (*O[a][b][\"seq\"],*O[a][b][\"mtx\"])\n",
        "\n",
        "            ##############################################\n",
        "            # filter to remove redundant sequences\n",
        "            ##############################################\n",
        "            ok = []\n",
        "            with open(\"tmp/tmp.fas\",\"w\") as fas_file:\n",
        "              fas_file.writelines([f\">{n}\\n{a+b}\\n\" for n,(a,b) in enumerate(zip(_seq_a,_seq_b))])\n",
        "            os.system(\"hhfilter -maxseq 1000000 -i tmp/tmp.fas -o tmp/tmp.id90.fas -id 90\")\n",
        "            for line in open(\"tmp/tmp.id90.fas\",\"r\"):\n",
        "              if line.startswith(\">\"): ok.append(int(line[1:]))\n",
        "            ##############################################            \n",
        "            print(f\"found {len(_seq_a)} pairs ({len(ok)} after filtering)\")\n",
        "\n",
        "            if len(_seq_a) > 0:\n",
        "              msa,mtx = [sequence],[[0]*len(sequence)]\n",
        "              for s_a,s_b,m_a,m_b in zip(_seq_a, _seq_b, _mtx_a, _mtx_b):\n",
        "                msa.append(_pad([a,b],[s_a,s_b],\"seq\"))\n",
        "                mtx.append(_pad([a,b],[m_a,m_b],\"mtx\"))\n",
        "              msas.append(msa)\n",
        "              deletion_matrices.append(mtx)\n",
        "    \n",
        "    '''\n",
        "    # triwise stitching (WIP)\n",
        "    if Ln > 2:\n",
        "      for a in range(Ln):\n",
        "        for b in range(a+1,Ln):\n",
        "          for c in range(b+1,Ln):\n",
        "            if O[a][b] is not None and O[b][c] is not None:\n",
        "              print(f\"attempting triwise stitch for {a} {b} {c}\")\n",
        "              list_ab = O[a][b][\"lab\"][1]\n",
        "              list_bc = O[b][c][\"lab\"][0]\n",
        "              msa,mtx = [sequence],[[0]*len(sequence)]\n",
        "              for i,l_b in enumerate(list_ab):\n",
        "                if l_b in list_bc:\n",
        "                  j = list_bc.index(l_b)\n",
        "                  s_a = O[a][b][\"seq\"][0][i]\n",
        "                  s_b = O[a][b][\"seq\"][1][i]\n",
        "                  s_c = O[b][c][\"seq\"][1][j]\n",
        "\n",
        "                  m_a = O[a][b][\"mtx\"][0][i]\n",
        "                  m_b = O[a][b][\"mtx\"][1][i]\n",
        "                  m_c = O[b][c][\"mtx\"][1][j]\n",
        "\n",
        "                  msa.append(_pad([a,b,c],[s_a,s_b,s_c],\"seq\"))\n",
        "                  mtx.append(_pad([a,b,c],[m_a,m_b,m_c],\"mtx\"))\n",
        "              if len(msa) > 1:\n",
        "                msas.append(msa)\n",
        "                deletion_matrices.append(mtx)\n",
        "                print(f\"found {len(msa)} triplets\")    \n",
        "    '''\n",
        "####################################################################################\n",
        "####################################################################################\n",
        "\n",
        "# save MSA as pickle\n",
        "pickle.dump({\"msas\":msas,\"deletion_matrices\":deletion_matrices},\n",
        "            open(os.path.join(output_dir,\"msa.pickle\"),\"wb\"))\n",
        "\n",
        "make_msa_plot = len(msas[0]) > 1\n",
        "if make_msa_plot:\n",
        "  plt = cf.plot_msas(msas, ori_sequence)\n",
        "  plt.savefig(os.path.join(output_dir,\"msa_coverage.png\"), bbox_inches = 'tight', dpi=200)\n",
        "  plt.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bQe3KeyTcv0n",
        "cellView": "form"
      },
      "source": [
        "#@title run alphafold\n",
        "num_relax = \"None\"\n",
        "rank_by = \"pLDDT\" #@param [\"pLDDT\",\"pTMscore\"]\n",
        "use_turbo = True #@param {type:\"boolean\"}\n",
        "max_msa = \"512:1024\" #@param [\"512:1024\", \"256:512\", \"128:256\", \"64:128\", \"32:64\"]\n",
        "max_msa_clusters, max_extra_msa = [int(x) for x in max_msa.split(\":\")]\n",
        "\n",
        "\n",
        "\n",
        "#@markdown - `rank_by` specify metric to use for ranking models (For protein-protein complexes, we recommend pTMscore)\n",
        "#@markdown - `use_turbo` introduces a few modifications (compile once, swap params, adjust max_msa) to speedup and reduce memory requirements. Disable for default behavior.\n",
        "#@markdown - `max_msa` defines: `max_msa_clusters:max_extra_msa` number of sequences to use. When adjusting after GPU crash, be sure to `Runtime` → `Restart runtime`. (Lowering will reduce GPU requirements, but may result in poor model quality. This option ignored if `use_turbo` is disabled)\n",
        "show_images = True #@param {type:\"boolean\"}\n",
        "#@markdown - `show_images` To make things more exciting we show images of the predicted structures as they are being generated. (WARNING: the order of images displayed does not reflect any ranking).\n",
        "#@markdown ---\n",
        "#@markdown #### Sampling options\n",
        "#@markdown There are two stochastic parts of the pipeline. Within the feature generation (choice of cluster centers) and within the model (dropout). \n",
        "#@markdown To get structure diversity, you can iterate through a fixed number of random_seeds (using `num_samples`) and/or enable dropout (using `is_training`).\n",
        "\n",
        "num_models = 5 #@param [1,2,3,4,5] {type:\"raw\"}\n",
        "use_ptm = True #@param {type:\"boolean\"}\n",
        "num_ensemble = 1 #@param [1,8] {type:\"raw\"}\n",
        "max_recycles = 3 #@param [1,3,6,12,24,48] {type:\"raw\"}\n",
        "tol = 0 #@param [0,0.1,0.5,1] {type:\"raw\"}\n",
        "is_training = False #@param {type:\"boolean\"}\n",
        "num_samples = 1 #@param [1,2,4,8,16,32] {type:\"raw\"}\n",
        "#@markdown - `num_models` specify how many model params to try. (5 recommended)\n",
        "#@markdown - `use_ptm` uses Deepmind's `ptm` finetuned model parameters to get PAE per structure. Disable to use the original model params. (Disabling may give alternative structures.)\n",
        "#@markdown - `num_ensemble` the trunk of the network is run multiple times with different random choices for the MSA cluster centers. (`1`=`default`, `8`=`casp14 setting`)\n",
        "#@markdown - `max_recycles` controls the maximum number of times the structure is fed back into the neural network for refinement. (3 recommended)\n",
        "#@markdown  - `tol` tolerance for deciding when to stop (CA-RMS between recycles)\n",
        "#@markdown - `is_training` enables the stochastic part of the model (dropout), when coupled with `num_samples` can be used to \"sample\" a diverse set of structures.\n",
        "#@markdown - `num_samples` number of random_seeds to try.\n",
        "subsample_msa = True #@param {type:\"boolean\"}\n",
        "#@markdown - `subsample_msa` subsample large MSA to `3E7/length` sequences to avoid crashing the preprocessing protocol. (This option ignored if `use_turbo` is disabled.)\n",
        "\n",
        "save_pae_json = True \n",
        "save_tmp_pdb = True\n",
        "\n",
        "\n",
        "if use_ptm == False and rank_by == \"pTMscore\":\n",
        "  print(\"WARNING: models will be ranked by pLDDT, 'use_ptm' is needed to compute pTMscore\")\n",
        "  rank_by = \"pLDDT\"\n",
        "\n",
        "#############################\n",
        "# delete old files\n",
        "#############################\n",
        "for f in os.listdir(output_dir):\n",
        "  if \"rank_\" in f:\n",
        "    os.remove(os.path.join(output_dir, f))\n",
        "\n",
        "#############################\n",
        "# homooligomerize\n",
        "#############################\n",
        "lengths = [len(seq) for seq in seqs]\n",
        "msas_mod, deletion_matrices_mod = cf.homooligomerize_heterooligomer(msas, deletion_matrices,\n",
        "                                                                    lengths, homooligomers)\n",
        "#############################\n",
        "# define input features\n",
        "#############################\n",
        "def _placeholder_template_feats(num_templates_, num_res_):\n",
        "  return {\n",
        "      'template_aatype': np.zeros([num_templates_, num_res_, 22], np.float32),\n",
        "      'template_all_atom_masks': np.zeros([num_templates_, num_res_, 37, 3], np.float32),\n",
        "      'template_all_atom_positions': np.zeros([num_templates_, num_res_, 37], np.float32),\n",
        "      'template_domain_names': np.zeros([num_templates_], np.float32),\n",
        "      'template_sum_probs': np.zeros([num_templates_], np.float32),\n",
        "  }\n",
        "\n",
        "num_res = len(full_sequence)\n",
        "feature_dict = {}\n",
        "feature_dict.update(pipeline.make_sequence_features(full_sequence, 'test', num_res))\n",
        "feature_dict.update(pipeline.make_msa_features(msas_mod, deletion_matrices=deletion_matrices_mod))\n",
        "if not use_turbo:\n",
        "  feature_dict.update(_placeholder_template_feats(0, num_res))\n",
        "\n",
        "def do_subsample_msa(F, random_seed=0):\n",
        "  '''subsample msa to avoid running out of memory'''\n",
        "  N = len(F[\"msa\"])\n",
        "  L = len(F[\"residue_index\"])\n",
        "  N_ = int(3E7/L)\n",
        "  if N > N_:\n",
        "    print(f\"whhhaaa... too many sequences ({N}) subsampling to {N_}\")\n",
        "    np.random.seed(random_seed)\n",
        "    idx = np.append(0,np.random.permutation(np.arange(1,N)))[:N_]\n",
        "    F_ = {}\n",
        "    F_[\"msa\"] = F[\"msa\"][idx]\n",
        "    F_[\"deletion_matrix_int\"] = F[\"deletion_matrix_int\"][idx]\n",
        "    F_[\"num_alignments\"] = np.full_like(F[\"num_alignments\"],N_)\n",
        "    for k in ['aatype', 'between_segment_residues',\n",
        "              'domain_name', 'residue_index',\n",
        "              'seq_length', 'sequence']:\n",
        "              F_[k] = F[k]\n",
        "    return F_\n",
        "  else:\n",
        "    return F\n",
        "\n",
        "################################\n",
        "# set chain breaks\n",
        "################################\n",
        "Ls = []\n",
        "for seq,h in zip(ori_sequence.split(\":\"),homooligomers):\n",
        "  Ls += [len(s) for s in seq.split(\"/\")] * h\n",
        "Ls_plot = sum([[len(seq)]*h for seq,h in zip(seqs,homooligomers)],[])\n",
        "feature_dict['residue_index'] = cf.chain_break(feature_dict['residue_index'], Ls)\n",
        "\n",
        "###########################\n",
        "# run alphafold\n",
        "###########################\n",
        "def parse_results(prediction_result, processed_feature_dict):\n",
        "  b_factors = prediction_result['plddt'][:,None] * prediction_result['structure_module']['final_atom_mask']  \n",
        "  dist_bins = jax.numpy.append(0,prediction_result[\"distogram\"][\"bin_edges\"])\n",
        "  dist_mtx = dist_bins[prediction_result[\"distogram\"][\"logits\"].argmax(-1)]\n",
        "  contact_mtx = jax.nn.softmax(prediction_result[\"distogram\"][\"logits\"])[:,:,dist_bins < 8].sum(-1)\n",
        "\n",
        "  out = {\"unrelaxed_protein\": protein.from_prediction(processed_feature_dict, prediction_result, b_factors=b_factors),\n",
        "         \"plddt\": prediction_result['plddt'],\n",
        "         \"pLDDT\": prediction_result['plddt'].mean(),\n",
        "         \"dists\": dist_mtx,\n",
        "         \"adj\": contact_mtx}\n",
        "\n",
        "  if \"ptm\" in prediction_result:\n",
        "    out.update({\"pae\": prediction_result['predicted_aligned_error'],\n",
        "                \"pTMscore\": prediction_result['ptm']})\n",
        "  return out\n",
        "\n",
        "model_names = ['model_1', 'model_2', 'model_3', 'model_4', 'model_5'][:num_models]\n",
        "total = len(model_names) * num_samples\n",
        "with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "  #######################################################################\n",
        "  # precompile model and recompile only if length changes\n",
        "  #######################################################################\n",
        "  if use_turbo:\n",
        "    name = \"model_5_ptm\" if use_ptm else \"model_5\"\n",
        "    N = len(feature_dict[\"msa\"])\n",
        "    L = len(feature_dict[\"residue_index\"])\n",
        "    compiled = (N, L, use_ptm, max_recycles, tol, num_ensemble, max_msa, is_training)\n",
        "    if \"COMPILED\" in dir():\n",
        "      if COMPILED != compiled: recompile = True\n",
        "    else: recompile = True\n",
        "    if recompile:\n",
        "      cf.clear_mem(\"gpu\")\n",
        "      cfg = config.model_config(name)      \n",
        "\n",
        "      # set size of msa (to reduce memory requirements)\n",
        "      msa_clusters = min(N, max_msa_clusters)\n",
        "      cfg.data.eval.max_msa_clusters = msa_clusters\n",
        "      cfg.data.common.max_extra_msa = max(min(N-msa_clusters,max_extra_msa),1)\n",
        "\n",
        "      cfg.data.common.num_recycle = max_recycles\n",
        "      cfg.model.num_recycle = max_recycles\n",
        "      cfg.model.recycle_tol = tol\n",
        "      cfg.data.eval.num_ensemble = num_ensemble\n",
        "\n",
        "      params = data.get_model_haiku_params(name,'./alphafold/data')\n",
        "      model_runner = model.RunModel(cfg, params, is_training=is_training)\n",
        "      COMPILED = compiled\n",
        "      recompile = False\n",
        "\n",
        "  else:\n",
        "    cf.clear_mem(\"gpu\")\n",
        "    recompile = True\n",
        "\n",
        "  # cleanup\n",
        "  if \"outs\" in dir(): del outs\n",
        "  outs = {}\n",
        "  cf.clear_mem(\"cpu\")  \n",
        "\n",
        "  #######################################################################\n",
        "  def report(key):\n",
        "    pbar.update(n=1)\n",
        "    o = outs[key]\n",
        "    line = f\"{key} recycles:{o['recycles']} tol:{o['tol']:.2f} pLDDT:{o['pLDDT']:.2f}\"\n",
        "    if use_ptm: line += f\" pTMscore:{o['pTMscore']:.2f}\"\n",
        "    print(line)\n",
        "    if show_images:\n",
        "      fig = cf.plot_protein(o['unrelaxed_protein'], Ls=Ls_plot, dpi=100)\n",
        "      plt.show()\n",
        "    if save_tmp_pdb:\n",
        "      tmp_pdb_path = os.path.join(output_dir,f'unranked_{key}_unrelaxed.pdb')\n",
        "      pdb_lines = protein.to_pdb(o['unrelaxed_protein'])\n",
        "      with open(tmp_pdb_path, 'w') as f: f.write(pdb_lines)\n",
        "\n",
        "  if use_turbo:\n",
        "    # go through each random_seed\n",
        "    for seed in range(num_samples):\n",
        "      \n",
        "      # prep input features\n",
        "      if subsample_msa:\n",
        "        sampled_feats_dict = do_subsample_msa(feature_dict, random_seed=seed)    \n",
        "        processed_feature_dict = model_runner.process_features(sampled_feats_dict, random_seed=seed)\n",
        "      else:\n",
        "        processed_feature_dict = model_runner.process_features(feature_dict, random_seed=seed)\n",
        "\n",
        "      # go through each model\n",
        "      for num, model_name in enumerate(model_names):\n",
        "        name = model_name+\"_ptm\" if use_ptm else model_name\n",
        "        key = f\"{name}_seed_{seed}\"\n",
        "        pbar.set_description(f'Running {key}')\n",
        "\n",
        "        # replace model parameters\n",
        "        params = data.get_model_haiku_params(name, './alphafold/data')\n",
        "        for k in model_runner.params.keys():\n",
        "          model_runner.params[k] = params[k]\n",
        "\n",
        "        # predict\n",
        "        prediction_result, (r, t) = cf.to(model_runner.predict(processed_feature_dict, random_seed=seed),\"cpu\")\n",
        "\n",
        "        # save results\n",
        "        outs[key] = parse_results(prediction_result, processed_feature_dict)\n",
        "        outs[key].update({\"recycles\":r, \"tol\":t})\n",
        "        report(key)\n",
        "\n",
        "        del prediction_result, params\n",
        "      del sampled_feats_dict, processed_feature_dict\n",
        "\n",
        "  else:  \n",
        "    # go through each model\n",
        "    for num, model_name in enumerate(model_names):\n",
        "      name = model_name+\"_ptm\" if use_ptm else model_name\n",
        "      params = data.get_model_haiku_params(name, './alphafold/data')  \n",
        "      cfg = config.model_config(name)\n",
        "      cfg.data.common.num_recycle = cfg.model.num_recycle = max_recycles\n",
        "      cfg.model.recycle_tol = tol\n",
        "      cfg.data.eval.num_ensemble = num_ensemble\n",
        "      model_runner = model.RunModel(cfg, params, is_training=is_training)\n",
        "\n",
        "      # go through each random_seed\n",
        "      for seed in range(num_samples):\n",
        "        key = f\"{name}_seed_{seed}\"\n",
        "        pbar.set_description(f'Running {key}')\n",
        "        processed_feature_dict = model_runner.process_features(feature_dict, random_seed=seed)\n",
        "        prediction_result, (r, t) = cf.to(model_runner.predict(processed_feature_dict, random_seed=seed),\"cpu\")\n",
        "        outs[key] = parse_results(prediction_result, processed_feature_dict)\n",
        "        outs[key].update({\"recycles\":r, \"tol\":t})\n",
        "        report(key)\n",
        "\n",
        "        # cleanup\n",
        "        del processed_feature_dict, prediction_result\n",
        "\n",
        "      del params, model_runner, cfg\n",
        "      cf.clear_mem(\"gpu\")\n",
        "\n",
        "  # delete old files\n",
        "  for f in os.listdir(output_dir):\n",
        "    if \"rank\" in f:\n",
        "      os.remove(os.path.join(output_dir, f))\n",
        "\n",
        "  # Find the best model according to the mean pLDDT.\n",
        "  model_rank = list(outs.keys())\n",
        "  model_rank = [model_rank[i] for i in np.argsort([outs[x][rank_by] for x in model_rank])[::-1]]\n",
        "\n",
        "  # Write out the prediction\n",
        "  for n,key in enumerate(model_rank):\n",
        "    prefix = f\"rank_{n+1}_{key}\" \n",
        "    pred_output_path = os.path.join(output_dir,f'{prefix}_unrelaxed.pdb')\n",
        "    fig = cf.plot_protein(outs[key][\"unrelaxed_protein\"], Ls=Ls_plot, dpi=200)\n",
        "    plt.savefig(os.path.join(output_dir,f'{prefix}.png'), bbox_inches = 'tight')\n",
        "    plt.close(fig)\n",
        "\n",
        "    pdb_lines = protein.to_pdb(outs[key][\"unrelaxed_protein\"])\n",
        "    with open(pred_output_path, 'w') as f:\n",
        "      f.write(pdb_lines)\n",
        "      \n",
        "############################################################\n",
        "print(f\"model rank based on {rank_by}\")\n",
        "for n,key in enumerate(model_rank):\n",
        "  print(f\"rank_{n+1}_{key} {rank_by}:{outs[key][rank_by]:.2f}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PnMog4kHUztt",
        "cellView": "form"
      },
      "source": [
        "#@title Refine structures with Amber-Relax (Optional)\n",
        "#@markdown If side-chain bond geometry is important to you, enable Amber-Relax by specifying how many top ranked structures you want relaxed. By default, we disable Amber-Relax since it barely moves the main-chain (backbone) structure and can overall double the runtime.\n",
        "num_relax = \"None\" #@param [\"None\", \"Top1\", \"Top5\", \"All\"] {type:\"string\"}\n",
        "if num_relax == \"None\":\n",
        "  num_relax = 0\n",
        "elif num_relax == \"Top1\":\n",
        "  num_relax = 1\n",
        "elif num_relax == \"Top5\":\n",
        "  num_relax = 5\n",
        "else:\n",
        "  num_relax = len(model_names) * num_samples\n",
        "\n",
        "#@markdown - `num_relax` specify how many of the top ranked structures to relax\n",
        "if num_relax > 0 and not os.path.isfile(\"stereo_chemical_props.txt\"):\n",
        "  try:\n",
        "    total = 45\n",
        "    with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      pbar.set_description(f'INSTALL AMBER')\n",
        "      with io.capture_output() as captured:\n",
        "        # Install OpenMM and pdbfixer.\n",
        "        %shell rm -rf /opt/conda\n",
        "        %shell wget -q -P /tmp \\\n",
        "          https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \\\n",
        "            && bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \\\n",
        "            && rm /tmp/Miniconda3-latest-Linux-x86_64.sh\n",
        "        pbar.update(4)\n",
        "\n",
        "        PATH=%env PATH\n",
        "        %env PATH=/opt/conda/bin:{PATH}\n",
        "        %shell conda update -qy conda \\\n",
        "            && conda install -qy -c conda-forge \\\n",
        "              python={python_version} \\\n",
        "              openmm=7.5.1 \\\n",
        "              pdbfixer\n",
        "        pbar.update(40)\n",
        "\n",
        "        %shell wget -q -P /content \\\n",
        "          https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
        "        pbar.update(1)\n",
        "        %shell mkdir -p /content/alphafold/common\n",
        "        %shell cp -f /content/stereo_chemical_props.txt /content/alphafold/common\n",
        "\n",
        "        # Apply OpenMM patch.\n",
        "        %shell pushd /opt/conda/lib/python{python_version}/site-packages/ && \\\n",
        "            patch -p0 < /content/alphafold/docker/openmm.patch && \\\n",
        "            popd\n",
        "\n",
        "  except subprocess.CalledProcessError:\n",
        "    print(captured)\n",
        "    raise\n",
        "\n",
        "if num_relax > 0:\n",
        "  if \"relax\" not in dir():\n",
        "    # add conda environment to path\n",
        "    sys.path.append(f\"/opt/conda/lib/python{python_version}/site-packages\")\n",
        "    \n",
        "    # import libraries\n",
        "    from alphafold.relax import relax\n",
        "    from alphafold.relax import utils\n",
        "\n",
        "  with tqdm.notebook.tqdm(total=num_relax, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "    pbar.set_description(f'AMBER relaxation')\n",
        "    for n,key in enumerate(model_rank):\n",
        "      if n < num_relax:\n",
        "        prefix = f\"rank_{n+1}_{key}\" \n",
        "        pred_output_path = os.path.join(output_dir,f'{prefix}_relaxed.pdb')\n",
        "        if not os.path.isfile(pred_output_path):\n",
        "          amber_relaxer = relax.AmberRelaxation(\n",
        "              max_iterations=0,\n",
        "              tolerance=2.39,\n",
        "              stiffness=10.0,\n",
        "              exclude_residues=[],\n",
        "              max_outer_iterations=20)\n",
        "          relaxed_pdb_lines, _, _ = amber_relaxer.process(prot=outs[key][\"unrelaxed_protein\"])        \n",
        "          with open(pred_output_path, 'w') as f:\n",
        "            f.write(relaxed_pdb_lines)\n",
        "        pbar.update(n=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KAZj6CBZTkJM",
        "cellView": "form"
      },
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "rank_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "\n",
        "key = model_rank[rank_num-1]\n",
        "prefix = f\"rank_{rank_num}_{key}\" \n",
        "pred_output_path = os.path.join(output_dir,f'{prefix}_relaxed.pdb')  \n",
        "if not os.path.isfile(pred_output_path):\n",
        "  pred_output_path = os.path.join(output_dir,f'{prefix}_unrelaxed.pdb') \n",
        "\n",
        "cf.show_pdb(pred_output_path, show_sidechains, show_mainchains, color, Ls=Ls_plot).show()\n",
        "if color == \"lDDT\": cf.plot_plddt_legend().show()  \n",
        "if use_ptm:\n",
        "  cf.plot_confidence(outs[key][\"plddt\"], outs[key][\"pae\"], Ls=Ls_plot).show()\n",
        "else:\n",
        "  cf.plot_confidence(outs[key][\"plddt\"], Ls=Ls_plot).show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XzK2Wve12GCk",
        "cellView": "form"
      },
      "source": [
        "#@title Extra outputs\n",
        "dpi =  100#@param {type:\"integer\"}\n",
        "save_to_txt = True #@param {type:\"boolean\"}\n",
        "save_pae_json = True #@param {type:\"boolean\"}\n",
        "#@markdown - save data used to generate contact and distogram plots below to text file (pae values can be found in json file if `use_ptm` is enabled)\n",
        "\n",
        "if use_ptm:\n",
        "  print(\"predicted alignment error\")\n",
        "  cf.plot_paes([outs[k][\"pae\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "  plt.savefig(os.path.join(output_dir,f'predicted_alignment_error.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "  plt.show()\n",
        "\n",
        "print(\"predicted contacts\")\n",
        "cf.plot_adjs([outs[k][\"adj\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "plt.savefig(os.path.join(output_dir,f'predicted_contacts.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "plt.show()\n",
        "\n",
        "print(\"predicted distogram\")\n",
        "cf.plot_dists([outs[k][\"dists\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "plt.savefig(os.path.join(output_dir,f'predicted_distogram.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "plt.show()\n",
        "\n",
        "print(\"predicted LDDT\")\n",
        "cf.plot_plddts([outs[k][\"plddt\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "plt.savefig(os.path.join(output_dir,f'predicted_LDDT.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "plt.show()\n",
        "\n",
        "def do_save_to_txt(filename, adj, dists):\n",
        "  adj = np.asarray(adj)\n",
        "  dists = np.asarray(dists)\n",
        "  L = len(adj)\n",
        "  with open(filename,\"w\") as out:\n",
        "    out.write(\"i\\tj\\taa_i\\taa_j\\tp(cbcb<8)\\tmaxdistbin\\n\")\n",
        "    for i in range(L):\n",
        "      for j in range(i+1,L):\n",
        "        if dists[i][j] < 21.68 or adj[i][j] >= 0.001:\n",
        "          line = f\"{i+1}\\t{j+1}\\t{full_sequence[i]}\\t{full_sequence[j]}\\t{adj[i][j]:.3f}\"\n",
        "          line += f\"\\t>{dists[i][j]:.2f}\" if dists[i][j] == 21.6875 else f\"\\t{dists[i][j]:.2f}\"\n",
        "          out.write(f\"{line}\\n\")\n",
        "\n",
        "for n,key in enumerate(model_rank):\n",
        "  if save_to_txt:\n",
        "    txt_filename = os.path.join(output_dir,f'rank_{n+1}_{key}.raw.txt')\n",
        "    do_save_to_txt(txt_filename,adj=outs[key][\"adj\"],dists=outs[key][\"dists\"])\n",
        "\n",
        "  if use_ptm and save_pae_json:\n",
        "    pae = outs[key][\"pae\"]\n",
        "    max_pae = pae.max()\n",
        "    # Save pLDDT and predicted aligned error (if it exists)\n",
        "    pae_output_path = os.path.join(output_dir,f'rank_{n+1}_{key}_pae.json')\n",
        "    # Save predicted aligned error in the same format as the AF EMBL DB\n",
        "    rounded_errors = np.round(np.asarray(pae), decimals=1)\n",
        "    indices = np.indices((len(rounded_errors), len(rounded_errors))) + 1\n",
        "    indices_1 = indices[0].flatten().tolist()\n",
        "    indices_2 = indices[1].flatten().tolist()\n",
        "    pae_data = json.dumps([{\n",
        "        'residue1': indices_1,\n",
        "        'residue2': indices_2,\n",
        "        'distance': rounded_errors.flatten().tolist(),\n",
        "        'max_predicted_aligned_error': max_pae.item()\n",
        "    }],\n",
        "                          indent=None,\n",
        "                          separators=(',', ':'))\n",
        "    with open(pae_output_path, 'w') as f:\n",
        "      f.write(pae_data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Riekgf0KQv_3",
        "cellView": "form"
      },
      "source": [
        "#@title Download prediction\n",
        "\n",
        "#@markdown Once this cell has been executed, a zip-archive with \n",
        "#@markdown the obtained prediction will be automatically downloaded \n",
        "#@markdown to your computer.\n",
        "\n",
        "# add settings file\n",
        "settings_path = os.path.join(output_dir,\"settings.txt\")\n",
        "with open(settings_path, \"w\") as text_file:\n",
        "  text_file.write(f\"notebook=https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/AlphaFold2_advanced.ipynb\\n\")\n",
        "  text_file.write(f\"sequence={ori_sequence}\\n\")\n",
        "  text_file.write(f\"msa_method={msa_method}\\n\")\n",
        "  if add_custom_msa:\n",
        "    text_file.write(f\"add_custom_msa={add_custom_msa} msa_format={msa_format}\\n\")\n",
        "  text_file.write(f\"homooligomer={homooligomer}\\n\")\n",
        "\n",
        "  text_file.write(f\"pair_mode={pair_mode}\\n\")\n",
        "  if pair_mode != \"unpaired\":\n",
        "    text_file.write(f\"pair_cov={pair_cov}\\n\")\n",
        "    text_file.write(f\"pair_qid={pair_qid}\\n\")\n",
        "\n",
        "  text_file.write(f\"max_msa={max_msa}\\n\")\n",
        "  text_file.write(f\"subsample_msa={subsample_msa}\\n\")  \n",
        "  text_file.write(f\"num_relax={num_relax}\\n\")\n",
        "  text_file.write(f\"use_turbo={use_turbo}\\n\")\n",
        "  text_file.write(f\"use_ptm={use_ptm}\\n\")\n",
        "  text_file.write(f\"rank_by={rank_by}\\n\")\n",
        "  text_file.write(f\"num_models={num_models}\\n\")\n",
        "  text_file.write(f\"num_samples={num_samples}\\n\")\n",
        "  text_file.write(f\"num_ensemble={num_ensemble}\\n\")\n",
        "  text_file.write(f\"max_recycles={max_recycles}\\n\")\n",
        "  text_file.write(f\"tol={tol}\\n\")\n",
        "  text_file.write(f\"is_training={is_training}\\n\")\n",
        "  text_file.write(f\"use_templates=False\\n\")\n",
        "  text_file.write(f\"-------------------------------------------------\\n\")\n",
        "\n",
        "  for n,key in enumerate(model_rank):\n",
        "    line = f\"rank_{n+1}_{key} pLDDT:{outs[key]['pLDDT']:.2f}\" + f\" pTMscore:{outs[key]['pTMscore']:.4f}\" if use_ptm else \"\"\n",
        "    text_file.write(line+\"\\n\")\n",
        "\n",
        "# --- Download the predictions ---\n",
        "!zip -FSr {output_dir}.zip {output_dir}\n",
        "files.download(f'{output_dir}.zip')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}